miχPODS: New baseline#

TODO

  • more intake usage

  • why does the stability diagram load multiple times

  • hvplot defaults

    • turn off scroll zoom

    • copy over presentation stuff

Catalog display#

Hide code cell source
%load_ext rich

# MOM6 run catalog
catalog = {
    "baseline": (
        "baseline",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.001.mixpods",
    ),
    "epbl": ("ePBL", "gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.epbl.001.mixpods"),
    "kpp.lmd.002": (
        "KPP Ri0=0.5",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.002.mixpods",
    ),
    "kpp.lmd.003": (
        "KPP Ri0=0.5, Ric=0.2,",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.003.mixpods",
    ),
    "kpp.lmd.004": (
        "KPP ν0=2.5, Ric=0.2, Ri0=0.5",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.004.mixpods",
    ),
    "baseline.N150": (
        "baseline N=150",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N150.baseline.mixpods",
    ),
    "kpp.lmd.004.N150": (
        "KPP ν0=2.5, Ric=0.2, Ri0=0.5, N=150",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N150.kpp.lmd.004.mixpods",
    ),
    "new_baseline.hb": (
        "KD=0, KV=0",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.hb",
    ),
    "new_baseline.kpp.lmd.004": (
        "KPP ν0=2.5, Ric=0.2, Ri0=0.5",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.kpp.lmd.004.mixpods",
    ),
    "new_baseline.kpp.lmd.005": (
        "KPP ν0=2.5, Ri0=0.5",
        "gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.kpp.lmd.005.mixpods",
    ),
}
catalog
The rich extension is already loaded. To reload it, use:
  %reload_ext rich
{
    'baseline': ('baseline', 'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.001.mixpods'),
    'epbl': ('ePBL', 'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.epbl.001.mixpods'),
    'kpp.lmd.002': (
        'KPP Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.002.mixpods'
    ),
    'kpp.lmd.003': (
        'KPP Ri0=0.5, Ric=0.2,',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.003.mixpods'
    ),
    'kpp.lmd.004': (
        'KPP ν0=2.5, Ric=0.2, Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.004.mixpods'
    ),
    'baseline.N150': ('baseline N=150', 'gmom.e23.GJRAv3.TL319_t061_zstar_N150.baseline.mixpods'),
    'kpp.lmd.004.N150': (
        'KPP ν0=2.5, Ric=0.2, Ri0=0.5, N=150',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N150.kpp.lmd.004.mixpods'
    ),
    'new_baseline.hb': ('KD=0, KV=0', 'gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.hb'),
    'new_baseline.kpp.lmd.004': (
        'KPP ν0=2.5, Ric=0.2, Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.kpp.lmd.004.mixpods'
    ),
    'new_baseline.kpp.lmd.005': (
        'KPP ν0=2.5, Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.kpp.lmd.005.mixpods'
    )
}

Setup#

Hide code cell source
%load_ext autoreload
%load_ext rich
%load_ext watermark

import warnings

import cf_xarray as cfxr
import dask_jobqueue
import distributed
import holoviews as hv
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tqdm
from datatree import DataTree

import xarray as xr

%aimport pump
from pump import mixpods

hv.notebook_extension("bokeh")

cfxr.set_options(warn_on_missing_variables=False)
xr.set_options(keep_attrs=True, display_expand_data=False)

plt.style.use("bmh")
plt.rcParams["figure.dpi"] = 140

mixing_layer_depth_criteria = {
    "boundary_layer_depth": {"name": "KPPhbl|KPP_OBLdepth|ePBL_h_ML"},
}

%watermark -iv
Hide code cell output
The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload
The rich extension is already loaded. To reload it, use:
  %reload_ext rich
The watermark extension is already loaded. To reload it, use:
  %reload_ext watermark
ipywidgets   : 8.0.6
cf_xarray    : 0.8.0
dask_jobqueue: 0.7.3
holoviews    : 1.15.4
matplotlib   : 3.7.1
distributed  : 2023.3.0
sys          : 3.10.10 | packaged by conda-forge | (main, Mar 24 2023, 20:08:06) [GCC 11.3.0]
xarray       : 2023.3.0
flox         : 0.6.10
pandas       : 1.5.3
dcpy         : 0.1.dev387+gd06c937
dask         : 2023.3.0
numpy        : 1.23.5
pump         : 1.0+247.g1f1c5e1.dirty
xrft         : 0.0.0
json         : 2.0.9
xgcm         : 0.6.1
tqdm         : 4.65.0
Hide code cell source
if "client" in locals():
    client.close()
    del client
if "cluster" in locals():
    cluster.close()

cluster = dask_jobqueue.PBSCluster(
    cores=4,  # The number of cores you want
    memory="23GB",  # Amount of memory
    processes=1,  # How many processes
    queue="casper",  # The type of queue to utilize (/glade/u/apps/dav/opt/usr/bin/execcasper)
    local_directory="/local_scratch/pbs.$PBS_JOBID/dask/spill",
    log_directory="/glade/scratch/dcherian/dask/",
    resource_spec="select=1:ncpus=4:mem=23GB",  # Specify resources
    project="ncgd0011",  # Input your project ID here
    walltime="02:00:00",  # Amount of wall time
    interface="ib0",  # Interface to use
)
cluster.adapt(minimum_jobs=1, maximum_jobs=4)
client = distributed.Client(cluster)
client
Hide code cell output

Client

Client-396a8bbb-d3d7-11ed-8a43-3cecef1ac722

Connection method: Cluster object Cluster type: dask_jobqueue.PBSCluster
Dashboard: https://jupyterhub.hpc.ucar.edu/stable/user/dcherian/proxy/8787/status

Cluster Info

Read#

Always read TAO#

tao_gridded = mixpods.load_tao()
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/xarray/core/dataset.py:255: UserWarning: The specified Dask chunks separate the stored chunks along dimension "depth" starting at index 42. This could degrade performance. Instead, consider rechunking after loading.
  warnings.warn(
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/xarray/core/dataset.py:255: UserWarning: The specified Dask chunks separate the stored chunks along dimension "time" starting at index 199726. This could degrade performance. Instead, consider rechunking after loading.
  warnings.warn(
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/xarray/core/dataset.py:255: UserWarning: The specified Dask chunks separate the stored chunks along dimension "longitude" starting at index 2. This could degrade performance. Instead, consider rechunking after loading.
  warnings.warn(

Subset Catalog, build tree#

catalog_sub = {
    k: catalog[k]
    for k in catalog.keys()
    if k == "kpp.lmd.004" or ("baseline" in k and "150" not in k)
}
display(catalog_sub)
{
    'baseline': ('baseline', 'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.001.mixpods'),
    'kpp.lmd.004': (
        'KPP ν0=2.5, Ric=0.2, Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.baseline.kpp.lmd.004.mixpods'
    ),
    'new_baseline.hb': ('KD=0, KV=0', 'gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.hb'),
    'new_baseline.kpp.lmd.004': (
        'KPP ν0=2.5, Ric=0.2, Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.kpp.lmd.004.mixpods'
    ),
    'new_baseline.kpp.lmd.005': (
        'KPP ν0=2.5, Ri0=0.5',
        'gmom.e23.GJRAv3.TL319_t061_zstar_N65.new_baseline.kpp.lmd.005.mixpods'
    )
}
datasets = {
    "TAO": tao_gridded,
    # "MITgcm": mitgcm,
}
with warnings.catch_warnings():
    warnings.simplefilter("ignore", category=UserWarning)
    warnings.simplefilter("ignore", category=FutureWarning)
    for short_name, (long_name, folder) in tqdm.tqdm(catalog_sub.items()):
        datasets.update(
            {
                short_name: mixpods.load_mom6_sections(folder).assign_attrs(
                    {"title": long_name}
                )
            }
        )
100%|██████████| 5/5 [06:23<00:00, 76.67s/it] 
datasets["les"] = les["0.0.-140.oct.average.month"].to_dataset()

# Offset LES to work with slicing below
datasets["les"]["time"] = datasets["les"]["time"] + pd.Timedelta(days=25 * 365)
tree = DataTree.from_dict(datasets)
tree
<xarray.DatasetView>
Dimensions:  ()
Data variables:
    *empty*

Modify catalog selection?#

if "les" in tree:
    tree["les"] = tree["les"].isel(z=slice(-2))
    tree["les"]["KT"].attrs["standard_name"] = "ocean_vertical_heat_diffusivity"

if "micro" in locals():
    tree.update(micro)

Post-process catalog subset#

tree = tree.sel(time=slice("2000", "2017"))
tree
<xarray.DatasetView>
Dimensions:  ()
Data variables:
    *empty*
ref = tree["TAO"].ds.reset_coords(drop=True).cf.sel(Z=slice(-120, None))
counts = np.minimum(ref["S2"].cf.count("Z"), ref["N2T"].cf.count("Z")).load()


def calc_histograms(ds):
    ds = ds.copy()
    ds["tao_mask"] = counts.reindex(time=ds.time, method="nearest") > 5
    ds["tao_mask"].attrs = {
        "description": "True when there are more than 5 5-m T, u, v in TAO dataset"
    }
    # ds = ds.where(ds.tao_mask)
    return ds.update(mixpods.pdf_N2S2(ds))


tree = tree.map_over_subtree(calc_histograms)

daily composite#

%autoreload

dailies = tree.map_over_subtree(mixpods.daily_composites)
dailies
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/numpy/lib/nanfunctions.py:1217: RuntimeWarning: All-NaN slice encountered
  r, k = function_base._ureduce(a, func=_nanmedian, axis=axis, out=out,
<xarray.DatasetView>
Dimensions:  ()
Data variables:
    *empty*

seasonal cycle#

tree = mixpods.persist_tree(tree)

Validate before continuing#

mixpods.validate_tree(tree)

TODO: Verify depth is normalized#

for name, ds in datasets.items():
    (ds.cf["sea_water_x_velocity"].cf["Z"].plot(marker=".", ls="none", label=name))
plt.legend()
<matplotlib.legend.Legend>
_images/2372712bb196933da3d938a29d3f04e0e217fe29d04bb036dbceb08ab9ce12da.png

Mean profiles#

%autoreload

S2 = mixpods.plot_profile_fill(tree, "S2", "S^2")
N2 = mixpods.plot_profile_fill(tree, "N2T", "N_T^2")
u = mixpods.plot_profile_fill(tree, "sea_water_x_velocity", "u")
T = mixpods.plot_profile_fill(tree, "sea_water_potential_temperature", "T")
Ri = mixpods.plot_median_Ri(tree)
(S2 + N2 + Ri + u + T).cols(3)

Stability Diagram#

%autoreload

mixpods.plot_stability_diagram_by_dataset(tree, nrows=2)
_images/67078def4b6e6abed5d4d3b9f4c9ced685ef78f679212ea748097dffb73b5b30.png

Daily composites#

%autoreload

mixpods.plot_daily_composites(dailies, ["eps", "KT"], logy=True, legend=False).opts(
    hv.opts.GridSpace(show_legend=True)
)
%autoreload

mixpods.plot_daily_composites(
    dailies, ["S2", "N2", "Rig_T"], logy=False, legend=False
).opts(hv.opts.GridSpace(show_legend=True, width=200), hv.opts.Overlay(frame_width=200))

Turbulence distributions#

handles = [
    mixpods.plot_distributions(tree, "chi", bins=np.linspace(-11, -4, 101), log=True),
    mixpods.plot_distributions(tree, "eps", bins=np.linspace(-11, -4, 101), log=True),
    mixpods.plot_distributions(
        tree, "ocean_vertical_heat_diffusivity", bins=np.linspace(-8, -1, 101), log=True
    ),
    # plot_distributions(tree, "Jq", bins=np.linspace(-1000, 0, 51), log=False),
    mixpods.plot_distributions(tree, "Rig_T", np.linspace(-0.5, 1.5, 61))
    * hv.VLine(0.25).opts(line_color="k"),
]
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/xarray/core/computation.py:769: RuntimeWarning: invalid value encountered in log10
  result_data = func(*input_data)
/glade/u/home/dcherian/miniconda3/envs/pump/lib/python3.10/site-packages/xarray/core/computation.py:769: RuntimeWarning: invalid value encountered in log10
  result_data = func(*input_data)
hv.Layout(handles).opts("Overlay", frame_width=600).cols(2)

Compare boundary layer depth#

hbl = (
    tree.drop_nodes("TAO")
    .dc.subset_nodes("KPP_OBLdepth")
    .dc.concatenate_nodes()
    .reset_coords(drop=True)
).load()
(
    # hbl.groupby("time.hour").mean().hvplot.line(by="node", flip_yaxis=True)
    hbl.groupby("time.hour").median().hvplot.line(by="node", flip_yaxis=True)
    + hbl.to_dataframe().hvplot.hist(
        by="node",
        bins=np.arange(0, 90, 1),
        normed=1,
        alpha=0.3,
        ylim=(0, 0.05),
        muted_alpha=0,
    )
).opts(hv.opts.Curve(invert_yaxis=True))